#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@File   : main.py
@Author : Liuli
@Date   : 2024-08-27 10:38 
"""
import numpy as np
import pandas as pd
from tqdm import tqdm
from asset_allocation import *
import backtrader as bt
from backtest import Strategy, AssetCommission, trade_list


if __name__ == "__main__":
    # backtrader 回测程序
    from WindPy import w
    w.start()

    start = '2016-01-01'
    end = '2024-10-31'
    start = pd.to_datetime(start).strftime('%Y%m%d')
    start_b = (pd.to_datetime(start) - pd.to_timedelta(365*2+60, 'days')).strftime('%Y%m%d')
    end = pd.to_datetime(end).strftime('%Y%m%d')
    # 300收益, 500收益, 中债-7-10年国债财富(总值)指数, 中债-3-5年期国债财富(总值)指数, 中债-0-1年国债财富(总值)指数,
    # 中债-优选投资级信用债指数, 美国全债市ETF-iShares, 标普500, SHFE黄金, SHFE铜, INE原油, 货币基金
    index = ['H00300.CSI', 'H00905.CSI', 'CBA06501.CS', 'CBA04601.CS', 'CBA14101.CS', 'CBA20901.CS', 'AGG.P',
             'SP500TR.SPI', 'AU.SHF', 'CU.SHF', 'B.IPE']  # 'SC.INE', 'H11025.CSI'
    price = w.wsd(index, "close", start_b, end, "")
    price = pd.DataFrame(np.array(price.Data).T, index=pd.to_datetime(price.Times), columns=price.Codes)
    price = price[pd.notna(price['H00300.CSI'])].copy()
    price.index.name = 'date'
    price.columns.name = 'code'
    ret = price.pct_change()
    # price = price.loc[start: end].copy()

    cerebro = bt.Cerebro()

    # 获取宏观因子数据用于计算权重
    from macro_factor import get_hidden_factor
    factor_data = get_hidden_factor(start_b, end)
    
    # 设置调仓周期（按月调仓）
    rebalance_dates = pd.date_range(start=start, end=end, freq='ME')

    # 生成调仓权重组合
    portfolio = {}
    for date in tqdm(rebalance_dates[-20:]):
        # 获取过去一年的收益率数据
        lookback_start = date - pd.DateOffset(years=2)
        hist_returns = ret.loc[lookback_start:date]
        
        # 获取对应期间的因子数据
        factor_period = factor_data.loc[lookback_start:date]
        
        # 计算风险平价权重
        try:
            weights = factor_risk_parity(hist_returns, factor_period)
            portfolio[date.strftime('%Y-%m-%d')] = pd.Series(weights, index=price.columns)
            if check_risk_parity_convergence(weights, hist_returns):
                print("优化结果收敛良好")
            else:
                print("优化可能未达到理想的风险平价状态")
        except Exception as e:
            print(f"Error calculating weights for {date}: {e}")
            break

    # 行情信息格式转换
    data = price.stack()
    data.name = 'open'
    data = data.reset_index(drop=False)
    data.rename(columns={'date': 'datetime', 'code': 'sec_code'}, inplace=True)
    data['datetime'] = pd.to_datetime(data['datetime'])
    data['close'] = data['open']
    data['high'] = data['open']
    data['low'] = data['open']
    data[['volume', 'openinterest']] = 0

    for code in tqdm(data['sec_code'].unique()):
        # 日期对齐
        date = pd.DataFrame(data['datetime'].unique(), columns=['datetime'])  # 获取回测区间内所有交易日
        data_ = data.query(f"sec_code=='{code}'")[
            ['datetime', 'open', 'high', 'low', 'close', 'volume', 'openinterest']]
        data_ = pd.merge(date, data_, how='left', on='datetime')
        data_ = data_.set_index("datetime")
        # 缺失值处理：日期对齐时会使得有些交易日的数据为空，所以需要对缺失数据进行填充
        data_.loc[:, ['open', 'high', 'low', 'close']] = data_.loc[:, ['open', 'high', 'low', 'close']].ffill()
        data_.loc[:, ['open', 'high', 'low', 'close']] = data_.loc[:, ['open', 'high', 'low', 'close']].fillna(0)
        datafeed = bt.feeds.PandasData(dataname=data_, fromdate=pd.to_datetime(data_.index.min()), 
                                       todate=pd.to_datetime(data_.index.max()))
        cerebro.adddata(datafeed, name=code)

    print("All Fund Data Done !")

    # 初始资金
    cerebro.broker.setcash(100000000.0)

    # 设置交易佣金
    cerebro.broker.addcommissioninfo(AssetCommission())

    # 当日下单，当日收盘价成交
    cerebro.broker.set_coc(True)

    # 添加策略，传入计算好的投资组合权重
    cerebro.addstrategy(Strategy, portfolio=portfolio)

    # 添加策略分析指标
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl') # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn') # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio') # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown') # 回撤
    cerebro.addanalyzer(trade_list, _name='tradelist') # 每笔交易盈亏情况
    # cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    # 添加观测器
    cerebro.addobserver(bt.observers.Broker)
    cerebro.addobserver(bt.observers.Trades)
    cerebro.addobserver(bt.observers.BuySell)
    cerebro.addobserver(bt.observers.DrawDown)
    cerebro.addobserver(bt.observers.TimeReturn)

    # 添加业绩基准时，需要事先将业绩基准的数据添加给 cerebro
    # 业绩基准使用二级债基
    bench = w.wsd("885007.WI", "close,open,high,low,volume", data['datetime'].min().strftime('%Y%m%d'), data['datetime'].max().strftime('%Y%m%d'), "")
    bench = pd.DataFrame(np.array(bench.Data).T, index=bench.Times, columns=bench.Fields)
    bench.index = pd.to_datetime(bench.index)
    bench = bench.reset_index(drop=False)

    bench.rename(columns={'index': 'datetime', 'OPEN': 'open', 'HIGH': 'high',
                          'LOW': 'low', 'CLOSE': 'close', 'VOLUME': 'volume'}, inplace=True)
    bench = bench.loc[(bench['datetime'] >= data['datetime'].min().strftime('%Y%m%d')) &
                              (bench['datetime'] <= data['datetime'].max().strftime('%Y%m%d'))].copy()
    bench['volume'] = 0
    bench['openinterest'] = 0
    bench['datetime'] = pd.to_datetime(bench['datetime'])
    bench = bench.set_index("datetime")
    bench = bt.feeds.PandasData(dataname=bench, fromdate=pd.to_datetime(bench.index.min()),
                                todate=pd.to_datetime(bench.index.max()))
    cerebro.adddata(bench, name='885001.WI')
    cerebro.addobserver(bt.observers.Benchmark, data=bench)

    # 启动回测
    result = cerebro.run(tradehistory=True)

    daily_return = pd.Series(result[0].analyzers.pnl.get_analysis())
    ret = pd.DataFrame(result[0].analyzers.tradelist.get_analysis())
    # 借助 pyfolio 进一步做回测结果分析
    # pyfolio = result[0].analyzers.pyfolio # 注意：后面不要调用 .get_analysis() 方法
    # 或者是 result[0].analyzers.getbyname('pyfolio')
    # returns, positions, transactions, gross_lev = pyfolio.get_pf_items()

    # import pyfolio as pf
    # pf.create_full_tear_sheet(returns)
    # 打印评价指标
    print("--------------- AnnualReturn -----------------")
    print(result[0].analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(result[0].analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(result[0].analyzers._DrawDown.get_analysis())
    # 可视化回测结果
    cerebro.plot(style='candlestick')
